import nltk
import pickle
from math import log
from torch.utils.data import Dataset, DataLoader
import pandas as pd
import numpy as np

from common.configs.path import paths
from common.configs.tools import label_map

from sklearn.model_selection import train_test_split

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter('./tfidf_nn_log')


if torch.cuda.is_available():
    print('gpu is available: {}'.format(torch.cuda.get_device_name(0)))
    print('device count: {}'.format(torch.cuda.device_count()))

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print('Using device:', device)
print()


def save_model(model, path):
    with open(path, 'wb') as f:
        pickle.dump(model, f)


def load_model(path):
    with open(path, 'rb') as f:
        model = pickle.load(f)
    return model


def load_npy(path):
    return np.load(path)


con_tfidf = load_npy(paths['output'] / 'tfidf/tfidf_con_3000.npy')
train_df = pd.read_csv(paths['train_data'])
y = train_df.label.apply(lambda e: label_map[e]).values
# y = y.reshape(y.shape[0], 1)


mask = np.logical_or(y == 10, y == 0, y == 14)

train_data = con_tfidf[:14009]
test_data = con_tfidf[14009:]

print(train_data.shape)
train_data = train_data[mask]
y = y[mask]
print(np.unique(y))
print(uu)
class_num = len(np.unique(y))
print(class_num)
# print(dd)

print(train_data.shape)


x, x_val, y, y_val = train_test_split(
    train_data, y, test_size=0.2, random_state=42)

sample_size = x.shape[0]
channels = 1
sequence_length = x.shape[1]

# x_train = x.reshape(-1, x.shape[1], 1).astype('float32')
x_train = x.reshape(sample_size, channels, sequence_length).astype('float32')

y_train = y


# x_val = x_val.reshape(-1, x_val.shape[1], 1).astype('float32')
x_val = x_val.reshape(x_val.shape[0], channels,
                      sequence_length).astype('float32')


# x_val = torch.from_numpy(x_val)
# y_val = torch.from_numpy(y_val)


class Data(Dataset):
    def __init__(self, x, y):
        self.x = torch.from_numpy(x).to(device)
        self.y = torch.from_numpy(y).to(device)
        self.len = self.x.shape[0]

    def __getitem__(self, index):
        return self.x[index], self.y[index]

    def __len__(self):
        return self.len


traindata_set = Data(x_train, y_train)
valdata_set = Data(x_val, y_val)

train_loader = DataLoader(dataset=traindata_set, batch_size=64)
test_loader = DataLoader(dataset=valdata_set, batch_size=64)


class NeuralNetwork(nn.Module):
    def __init__(self):
        super(NeuralNetwork, self).__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Conv1d(in_channels=1, out_channels=16, kernel_size=3),
            nn.MaxPool1d(4,),
            # nn.Conv1d(in_channels=16, out_channels=16, kernel_size=3),
            # nn.MaxPool1d(4,),

            nn.Linear(749, 1024),
            nn.Flatten(),
            nn.ReLU(),
            nn.Linear(1024, 1024),
            nn.ReLU(),
            nn.Linear(1024, 512),
            nn.ReLU(),
            nn.Linear(512, class_num),
        )

    def forward(self, x):
        # x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        # print('!!!')
        # logits = self.flatten(logits)

        # print('shape logits', logits.shape)
        return logits


def train_loop(dataloader, model, loss_fn, optimizer):
    size = len(dataloader.dataset)
    for batch, (X, y) in enumerate(dataloader):
        # Compute prediction and loss
        pred = model(X)
        loss = loss_fn(pred, y)

        # Backpropagation
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if batch % 100 == 0:
            loss, current = loss.item(), batch * len(X)
            print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")


def test_loop(dataloader, model, loss_fn):
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    test_loss, correct = 0, 0

    with torch.no_grad():
        for X, y in dataloader:
            pred = model(X)
            print(pred.shape)
            test_loss += loss_fn(pred, y).item()
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()
            # print(pred)

    test_loss /= num_batches
    correct /= size
    print(
        f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")
    return test_loss


model = NeuralNetwork().to(device)
print(model)
learning_rate = 0.005

loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)


epochs = 25
for t in range(epochs):
    print(f"Epoch {t+1}\n-------------------------------")
    train_loop(train_loader, model, loss_fn, optimizer)
    eval = test_loop(test_loader, model, loss_fn)
    writer.add_scalar('cnn', eval, t)

print("Done!")
